from agno.models.openai import OpenAILike

from configs.basic_config import ALI_GENERATION_MODEL, AGENT_URL, ALI_DASHSCOPE_API_KEY
from agno.agent import Agent
from agno.db.sqlite import SqliteDb
from agno.workflow.step import Step
from agno.workflow.workflow import Workflow
# --- 修正 RunContext 导入 ---
from agno.run import RunContext # 假设 RunContext 在 agno.run 中

db = SqliteDb(db_file="tmp/workflow.db")


# 定义工具以在工作流会话状态中管理购物清单
def add_item(run_context: RunContext, item: str) -> str:
    """将项目添加到工作流会话状态中的购物清单。

    Args:
        item (str): 要添加到购物清单的项目
    """
    if not run_context.session_state:
        run_context.session_state = {}

    # 检查项目是否已存在 (不区分大小写)
    existing_items = [
        existing_item.lower() for existing_item in run_context.session_state["shopping_list"]
    ]
    if item.lower() not in existing_items:
        run_context.session_state["shopping_list"].append(item)
        return f"已将 '{item}' 添加到购物清单。"
    else:
        return f"'{item}' 已在购物清单中。"


def remove_item(run_context: RunContext, item: str) -> str:
    """从工作流会话状态中的购物清单中移除项目。

    Args:
        item (str): 要从购物清单中移除的项目
    """
    if not run_context.session_state:
        run_context.session_state = {}

    if len(run_context.session_state["shopping_list"]) == 0:
        return f"购物清单为空。无法移除 '{item}'。"

    # 查找并移除项目 (不区分大小写)
    shopping_list = run_context.session_state["shopping_list"]
    for i, existing_item in enumerate(shopping_list):
        if existing_item.lower() == item.lower():
            removed_item = shopping_list.pop(i)
            return f"已从购物清单中移除 '{removed_item}'。"

    return f"未在购物清单中找到 '{item}'。"


def remove_all_items(run_context: RunContext) -> str:
    """从工作流会话状态中的购物清单中移除所有项目。"""
    if not run_context.session_state:
        run_context.session_state = {}

    run_context.session_state["shopping_list"] = []
    return "已从购物清单中移除所有项目。"


def list_items(run_context: RunContext) -> str:
    """列出工作流会话状态中的所有购物清单项目。"""
    if not run_context.session_state:
        run_context.session_state = {}

    if len(run_context.session_state["shopping_list"]) == 0:
        return "购物清单为空。"

    items = run_context.session_state["shopping_list"]
    items_str = "\n".join([f"- {item}" for item in items])
    return f"购物清单:\n{items_str}"


# 创建使用利用工作流会话状态的工具的代理
shopping_assistant = Agent(
    name="购物助手",
    # --- 修改模型参数 ---
    model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY, base_url=AGENT_URL),
    tools=[add_item, remove_item, list_items],
    instructions=[
        "你是一个有用的购物助手。",
        "你可以通过添加、删除和列出项目来帮助用户管理他们的购物清单。",
        "总是使用提供的工具与购物清单进行交互。",
        "在你的回复中要友好和乐于助人。",
    ],
)

list_manager = Agent(
    name="清单管理员",
    # --- 修改模型参数 ---
    model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY, base_url=AGENT_URL),
    tools=[list_items, remove_all_items],
    instructions=[
        "你是一个清单管理专家。",
        "你可以查看当前的购物清单并在需要时清除它。",
        "当被要求时总是显示当前清单。",
        "清楚地向用户确认操作。",
    ],
)

# 创建步骤
manage_items_step = Step(
    name="manage_items",
    description="帮助管理购物清单项目 (添加/删除)",
    agent=shopping_assistant,
)

view_list_step = Step(
    name="view_list",
    description="查看和管理完整的购物清单",
    agent=list_manager,
)

# 创建带有 workflow_session_state 的工作流
shopping_workflow = Workflow(
    name="购物清单工作流",
    #db=db,
    steps=[manage_items_step, view_list_step],
    session_state={"shopping_list": []},
)

if __name__ == "__main__":
    # 示例 1: 向购物清单添加项目
    print("=== 示例 1: 添加项目 ===")
    shopping_workflow.print_response(
        input="请将牛奶、面包和鸡蛋添加到我的购物清单。"
    )
    print("工作流会话状态:", shopping_workflow.get_session_state())

    # 示例 2: 添加更多项目并查看清单
    print("\n=== 示例 2: 添加更多项目 ===")
    shopping_workflow.print_response(
        input="将苹果和香蕉添加到清单中，然后告诉我完整的清单。"
    )
    print("工作流会话状态:", shopping_workflow.get_session_state())

    # 示例 3: 移除项目
    print("\n=== 示例 3: 移除项目 ===")
    shopping_workflow.print_response(
        input="从清单中移除面包，然后告诉我剩下的东西。"
    )
    print("工作流会话状态:", shopping_workflow.get_session_state())

    # 示例 4: 清空整个清单
    print("\n=== 示例 4: 清空清单 ===")
    shopping_workflow.print_response(
        input="清空整个购物清单，并确认它已为空。"
    )
    print("最终工作流会话状态:", shopping_workflow.get_session_state())